"""
  Copyright FunASR (https://github.com/alibaba-damo-academy/FunASR). All Rights
  Reserved. MIT License  (https://opensource.org/licenses/MIT)
  
  2023-2024 by zhaomingwork@qq.com  
"""

# pip install websocket-client
# apt install ffmpeg 

import time
import traceback
from biz.integrations.funasr.funasr_core import FunasrCore

# class for recognizer in websocket
class FunasrApi:
    """
    python asr recognizer lib
    """
    def __init__(
        self,
        uri="wss://www.funasr.com:10096/",
        timeout=1000,
        msg_callback=None,
        max_retry=3,
        keep_alive=False
    ):
        """
        uri: ws or wss server uri
        msg_callback: for message received
        timeout: timeout for get result
        """
        try:
            self.uri=uri
            self.timeout=timeout
            self.msg_callback=msg_callback
            self.funasr_core=None
            self.max_retry = max_retry
            self.keep_alive = keep_alive
        except Exception as e:
            print("Exception:", e)
            traceback.print_exc()

    def ensure_connection(self):
        """确保连接存在，如果不存在或已断开则重新连接"""
        retry_count = 0
        while retry_count < self.max_retry:
            try:
                if self.funasr_core is None or not self.funasr_core.is_connected():
                    if self.funasr_core is not None:
                        try:
                            self.funasr_core.close()
                        except:
                            pass
                    self.funasr_core = FunasrCore(self.uri, msg_callback=self.msg_callback, timeout=self.timeout)
                    self.funasr_core.new_connection()
                return True
            except Exception as e:
                retry_count += 1
                print(f"连接失败，正在重试 ({retry_count}/{self.max_retry}): {e}")
                time.sleep(1)  # 重试前等待1秒
        print(f"连接失败，已达到最大重试次数 {self.max_retry}")
        return False

    def new_core(self,msg_callback=None):
     try:
         if self.funasr_core is not None:
            self.funasr_core.close()
            
         if msg_callback==None:
            msg_callback=self.msg_callback
         funasr_core=FunasrCore(self.uri,msg_callback=msg_callback,timeout=self.timeout)
         funasr_core.new_connection()
         self.funasr_core=funasr_core
         return funasr_core
         
     except Exception as e:
            print("init_core",e)
            exit(0)
    
    def rec_buf(self,audio_buf,ffmpeg_decode=False):
        """识别音频缓冲区"""
        try:
            # 使用现有连接或确保连接可用
            if self.keep_alive:
                if not self.ensure_connection():
                    raise Exception("无法建立连接")
                # 使用现有连接
                self.funasr_core.rec_buf(audio_buf, ffmpeg_decode=ffmpeg_decode)
                result = self.funasr_core.get_result_realtime()
                if self.funasr_core is not None:
                    self.funasr_core.is_final = False
                    self.funasr_core.rec_text = ""
                    self.funasr_core.rec_file_len = 0
                return result
            else:
                funasr_core = self.new_core()
                funasr_core.rec_buf(audio_buf, ffmpeg_decode=ffmpeg_decode)
                return funasr_core.get_result()
        except Exception as e:
            print("rec_buf", e)
            # 连接失败时尝试重新连接
            if self.keep_alive:
                print("尝试重新连接...")
                self.ensure_connection()
            return None


 
    def funasr_close(self):
        """关闭连接"""
        try:
            if self.funasr_core is not None:
                self.funasr_core.close()
        except Exception as e:
            print("Exception:", e)
            traceback.print_exc()
